import numpy as np
import matplotlib.pyplot as plt
import math
import mpl_toolkits.axisartist as axisartist


def func1(t, T, delta):
    # 方波信号
    if (t + delta) % T <= 2 * delta:
        return 1
    else:
        return 0

def func2(f, delta):
    # 方波信号的双边谱
    f = abs(f)
    return math.sin(2 * math.pi * f * delta) / (math.pi * f)

def func3(f, delta):
    # 方波信号的单边谱
    if f < 0.01:
        return 0
    elif abs(f) <= 0.01:
        return func2(f, delta)
    else:
        return 2 * func2(f, delta)

def func4(t, T, delta):
    # 三角信号
    if (t + delta) % T <= delta:
        return (t + delta) % T / delta
    elif (t + delta) % T <= 2 * delta:
        return 1 - t % T / delta
    else:
        return 0

def func5(f, delta):
    # 三角信号的双边谱
    return (math.sin(math.pi * f * delta)) ** 2 / (math.pi * f) ** 2 / delta

def func6(f, delta):
    # 三角信号的单边谱
    if f < 0.01:
        return 0
    elif abs(f) <= 0.01:
        return func5(f, delta)
    else:
        return 2 * func5(f, delta)

def set_axis(fig, number, title, x_range=(-10, 15), y_range=(-1, 1)):
    ax = axisartist.Subplot(fig, number)
    fig.add_axes(ax)
    ax.axis[:].set_visible(False)
    ax.axis["x"] = ax.new_floating_axis(0,0)
    ax.axis["x"].set_axisline_style("->", size = 1.0)
    ax.axis["y"] = ax.new_floating_axis(1,0)
    ax.axis["y"].set_axisline_style("->", size = 1.0)
    ax.axis["x"].set_axis_direction("bottom")
    ax.axis["y"].set_axis_direction("left")
    ax.set_xlim(x_range[0], x_range[1])
    ax.set_ylim(y_range[0], y_range[1])
    ax.set_title(title)

def mod(x, f):
    x = abs(x)
    return x - math.floor(x / f) * f < 1e-2


if __name__ == "__main__":

    fig = plt.figure(figsize=(10, 7))
    plt.rcParams['font.family'] = ['Arial Unicode MS']

    # 绘制方波信号
    T = 6
    delta = 0.5
    t = np.arange(-2 * T, 2 * T, 0.01)
    f = np.arange(-2 * T, 2 * T, 0.01)
    n_f = np.array([i for i in f if mod(i, 1 / T)])

    f1 = np.array([func1(i, T, delta) for i in t])
    set_axis(fig, 231, "方波信号", (-9, 9), (-1, 2))
    plt.plot(t, f1)

    # 绘制方波信号的复数形式频谱
    f2 = np.array([1 / T * func2(i, delta) for i in n_f])
    set_axis(fig, 232, "方波信号的复数形式频谱", (-5, 5), (min(f2) * 1.2, max(f2) * 1.2))
    plt.stem(n_f, f2, markerfmt='.')

    # 绘制方波信号的三角形式频谱
    f3 = np.array([1 / T * func3(i, delta) for i in n_f])
    set_axis(fig, 233, "方波信号三角形式频谱", (0, 7), (min(f3) * 1.2, max(f3) * 1.2))
    plt.stem(n_f, f3, markerfmt='.')

    # 绘制三角信号
    T = 6
    delta = 0.5
    t = np.arange(-2 * T, 2 * T, 0.01)
    t = np.arange(-2 * T, 2 * T, 0.01)
    n_f = np.array([i for i in f if mod(i, 1 / T)])

    f4 = np.array([func4(i, T, delta) for i in t])
    set_axis(fig, 234, "三角信号", (-19, 19), (-1, 2))
    plt.plot(t, f4)

    # 绘制三角信号的复数形式频谱
    f5 = np.array([1 / T * func5(i, delta) for i in n_f])
    set_axis(fig, 235, "三角信号的复数形式频谱", (-5, 5), (min(f5) * 0.8, max(f5) * 1.2))
    plt.stem(n_f, f5, markerfmt='.')

    # 绘制三角信号的三角形式频谱
    f6 = np.array([1 / T * func6(i, delta) for i in n_f])
    set_axis(fig, 236, "三角信号的三角形式频谱", (0, 7), (min(f6) * 0.8, max(f6) * 1.2))
    plt.stem(n_f, f6, markerfmt='.')
    
    plt.show()